import glob
import json
import shutil

import cv2

import os


def t1():
    ls = glob.glob(r'D:\DATA\20250611HKBZ\caitu\20250613_zxc_labelme0_164\*.jpg')
    for pa in ls:
        filename = os.path.basename(pa)  # 1.jpg
        parent_folder = os.path.dirname(pa)
        ind_str = filename.split('.')[0]
        new_path = f'{parent_folder}/{ind_str.zfill(6)}.jpg'  # 000001.jpg
        os.rename(pa, new_path)
        print(f'{pa} to {new_path}')


def add_clsinfor2filename():
    img_annotation_dir = r'D:\DATA\20250611HKBZ\caitu\testdata170x'

    cls_names = ['aircraft', 'cargo_door', 'chock', 'refueling_truck', 'lifting_platform',
                 'transfer_car', 'staff', 'connecting', 'catering_truck', 'belt_loader',
                 'power_cable', 'pre_conditioning_air', 'thumbs_up', 'ULD', 'bag', 'push_back_tractor',
                 'no_detection']

    img_ls = glob.glob(f'{img_annotation_dir}/*.jpg')
    for img_path in img_ls:
        json_path = img_path.replace('.jpg', '.json')
        json_dict = json.load(open(json_path))
        cls_count_dict = {}
        for rect_shape_dict in json_dict['shapes']:
            cls_name = rect_shape_dict['label']
            if cls_name not in cls_count_dict:
                cls_count_dict[cls_name] = 0
            cls_count_dict[cls_name] += 1
        cs = list(cls_count_dict.keys())
        img_new_filename = os.path.basename(img_path)
        for cls_name in cls_names:
            if cls_name in cs:
                img_new_filename = f'{cls_name}_{img_new_filename}'

        json_new_filename = img_new_filename.replace('.jpg', '.json')
        json_new_path = f'{img_annotation_dir}/{json_new_filename}'
        img_new_path = f'{img_annotation_dir}/{img_new_filename}'
        shutil.move(json_path, json_new_path)
        shutil.move(img_path, img_new_path)

        print(f'{json_path} to {json_new_path}')
        print(f'{img_path} to {img_new_path}')


def frames2video(frames_glob, video_save_path):
    # frames_glob = r'D:\DATA\20250611HKBZ\caitu\20250619\action\3\DengJiQiao1JiuWei.mp4_down20_frame_res\*.jpg'
    # video_save_path = r'D:\DATA\20250611HKBZ\caitu\20250619\action\3\DengJiQiao1JiuWei_res.mp4'
    fps = 2  # min 1
    # w, h = 2160, 3840
    w, h = 3840, 2160
    mp4 = cv2.VideoWriter_fourcc(*'mp4v')  # MJPG  XVID
    res = cv2.VideoWriter(video_save_path, mp4, fps, (w, h), True)  # WH

    frame_paths = glob.glob(frames_glob)
    print(f'fps: {fps}, total: {len(frame_paths)}, w: {w}, h: {h}')

    for ind, frame_path in enumerate(frame_paths):
        print(f'{ind}/{len(frame_paths)} {frame_path}')
        frame = cv2.imread(frame_path)
        frame_ind = os.path.basename(frame_path)
        cv2.putText(frame, frame_ind, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        res.write(frame)
    res.release()
    print('Done.')


if __name__ == '__main__':
    # t1()
    # add_clsinfor2filename()
    frames2video(
        frames_glob=r'D:\DATA\20250611HKBZ\caitu\20250619\action\3\DengJiQiao1JiuWei.mp4_down20_frame_res\*.jpg',
        video_save_path=r'D:\DATA\20250611HKBZ\caitu\20250619\action\3\DengJiQiao1JiuWei_res.mp4')
